% calculates welfare under epstein-zin preferences
% for HHW contract and PHI contract

function [EZ_welfare_HHW V]= EZ_welfare_HHW_v4(Wi,CAPH,MU,maxreps,tolFP,rho,psi,alpha,invG,G,prob_start)

T = size(Wi,1);

Nstates = size(MU,1);

% consumption will never decrease

Y1 = Wi(:,1);

Cg1 = HHW_cg(Y1,MU,CAPH,Nstates,rho,T,tolFP,maxreps);

C1 = Cg1(:,1);

minC= floor(min(C1(C1>=0,1)));

%minC= 0;

% grid should be fine enough
% let's say 100 points between min consumption and the max W?
% i will assume no one starts in the states that would yield negative
% consumption so i avoid issues with the CES in the EZ preferences 


maxC = max(Wi)-min(min(MU));
delta = maxC-minC;
step = round(delta/1000);
step = max(step,1);
c_grid = [minC:step:maxC];
C_grid = repmat(c_grid,Nstates,1);

% mortality

C_grid(Nstates,:)=0;
ngrid = size(C_grid,2);
%ngrid
   
C_bar_f=Cg1(:,T);

%maybe unnecesary?

%C_bar_f(Nstates,1)=0;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% period T-1
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

t  = T-1;
  
% C_f = C{t+1}|c_t, lambda_{t+1}
% consumption grid(cols) and future lambda in rows

% future consumption in HHW contract
% lapsation occurs iff current c is lower than consumption 
% warranty
C_f = C_grid.*(C_grid>repmat(C_bar_f,1,ngrid))+repmat(C_bar_f,1,ngrid).*(C_grid<=repmat(C_bar_f,1,ngrid));

% C_f(Nstates,:) = 0;
% The value function in each future lambda (lambda_t+1)
% For terminal period is just (1-rho)*C_f

V_f = (1-rho)^(1/(1-1/psi))*C_f;

%V_f = CES_agg(C_f,0,psi,rho);

% Take expectations over lambda 

capt = CAPH(:,:,t);
% REPLACE
GV_f = G(V_f,alpha);
%GV_f = G_cara(V_f,gamma);
GV_f(Nstates,:)=0;
EGV_f = capt*GV_f;  

%inverse operator
  
% REPLACE
invGEGV_f = invG(EGV_f,alpha); 
%invGEGV_f = invG_cara(EGV_f,gamma); 
invGEGV_f(Nstates,:)=0;

% for our given p; and for each state; what is V
% using CES aggregator
V = CES_agg(C_grid,invGEGV_f,psi,rho);


t= T-2;

while (t>=1)

  capt = CAPH(:,:,t);  
  
  C_bar_f=Cg1(:,t+1);
    

  %% for every possible consumption/lambda in the future; 
  %% this is V in the future

  V_f = V;

  % If you land in the state in the row (j); and 
  % you currently have consumption (i)
  % your future consumption is C_f(i,j)
  % so this is C_{t+1} | lambda_{t+1}, c_t

  C_f = C_grid.*(C_grid>repmat(C_bar_f,1,ngrid))+...
    repmat(C_bar_f,1,ngrid).*(C_grid<=repmat(C_bar_f,1,ngrid));

  % just a fix so that it does find an index
  C_f(Nstates,:) = minC;
  
  % round to the next integer
  % this makes the whole thing an approximation
  
  C_fmCmin = C_f-minC;
  C_fmCmin = round(C_fmCmin/step)*step;
  %C_f = round(C_f/step)*step;
  Col_C_f = C_fmCmin/step+1;

  %now I need to read V(C_{t+1}}|lambda_{t+1},c_t}  
  % if future state is in row r
  % and future consumption is equal to C
  % then I need to read Vf at: 
  % row r and 
  % the column that corresponds to C in the C_grid 
  % which means I need to read the following index  
    
  index =(Col_C_f-1)*Nstates+repmat([1:Nstates]',1,ngrid);
  %max(max(index))
  
  GV_f = G(V_f(index),alpha);
  GV_f(Nstates,:)=0;
  EGV_f = capt*GV_f;  
  invGEGV_f = invG(EGV_f,alpha); 
%  invGEGV_f = invG_cara(EGV_f,gamma); 

  invGEGV_f(Nstates,:)=0;
  V = CES_agg(C_grid,invGEGV_f,psi,rho);
   
  t = t-1;  

end 

  %starting consumption
  
  %C1 = Cg1(:,1);
  
  
  
  C1(Nstates,1) = 0;
  delta = (repmat(c_grid,Nstates,1) -repmat(C1,1,ngrid)).^2;
  mindelta = min(delta,[],2);  
  C1_loc = delta==repmat(mindelta,1,ngrid);
  V1 = sum(V.*C1_loc,2);

 % CERTAINTY EQUIVALENT
  % TAKING INTO ACCOUNT THE FIRST-PERIOD UNCERTAINTY
  
  CAPC =eye(length(MU(:,:,1))); 
  psurv = zeros(T,1);
  psurv(1,1) = 1;
  d = zeros(T,1);
  d(1,1)=1;

  for t = 2:T
      CAPC = CAPC*CAPH(:,:,t-1);
      psurv(t,1) = 1-CAPC(1,length(MU(:,:,1)));
      d(t,1) = psurv(t,1)*rho^(t-1);
  end

  
  GV1 = G(V1,alpha);
  EGV1 = prob_start*GV1;
  V0 = invG(EGV1,alpha);
  a = 1/(1-rho)*V0^(1-1/psi)/sum(d);
  EZ_welfare_HHW= a^(1/(1-1/psi));

 % a = G(V0,alpha)/(1-rho)/sum(d);
 % EZ_welfare_HHW_v3= invG(a,alpha);

  
end